from safetensors.torch import load_file
def tensor_reader(file_path):
    try:
        tensors = load_file(file_path)
        tensor_info = {name: tensor.shape for name, tensor in tensors.items()}
        return tensor_info
    except Exception as e:
        return f'Error reading tensor: {e}', {}
import os
files = [file for file in os.listdir() if file.endswith('.safetensors')]
if files:
    print('Available .safetensors files:')
    for idx, file in enumerate(files, 1):
        print(f'{idx}. {file}')
    choice = input('Enter the number of the file to read tensor info from: ')
    try:
        choice = int(choice)
        if 1 <= choice <= len(files):
            file_path = files[choice - 1]
            tensor_info = tensor_reader(file_path)
            print('Tensors:')
            for name, shape in tensor_info.items():
                print(f'  {name}: {shape}')
        else:
            print('Invalid selection.')
    except (ValueError, IndexError):
        print('Invalid choice. Please enter a valid number.')
else:
    print('No safetensors files are available in the current directory.')
    input('--- Press ENTER To Exit ---')